"""
Coze集成模块
整合OpenAPI客户端、WebSocket客户端和配置管理器
"""

import asyncio
import logging
from typing import Optional, Dict, Any, Union, Callable
from dataclasses import dataclass
from datetime import datetime
import numpy as np

from .coze_openapi_client import (
    CozeOpenAPIClient, 
    CozeAPIConfig, 
    ChatResponse,
    CozeAPIError,
    CozeAuthenticationError,
    CozeRateLimitError,
    CozeNetworkError
)
from .coze_websocket_client import (
    CozeWebSocketClient,
    WebSocketConfig,
    VisionWebSocketEventHandler
)
from .coze_config_manager import (
    CozeConfigManager,
    CozeAPICredentials,
    CozeAgentConfig,
    CozeConnectionConfig,
    CozeFlowConfig
)

logger = logging.getLogger(__name__)


@dataclass
class VisionQueryRequest:
    """视觉查询请求"""
    query_id: str
    user_query: str
    image_data: bytes
    audio_data: Optional[bytes] = None
    use_websocket: bool = False
    timeout: int = 30


@dataclass
class VisionQueryResponse:
    """视觉查询响应"""
    query_id: str
    response_text: str
    confidence_score: float
    processing_time_ms: int
    model_used: str
    method_used: str  # "openapi" | "websocket"
    timestamp: datetime
    error: Optional[str] = None


class CozeIntegrationError(Exception):
    """Coze集成错误"""
    pass


class CozeIntegration:
    """
    Coze集成管理器
    统一管理OpenAPI和WebSocket客户端
    """
    
    def __init__(self, config_dir: Optional[str] = None):
        """
        初始化Coze集成管理器
        
        Args:
            config_dir: 配置文件目录
        """
        self.config_manager = CozeConfigManager(config_dir)
        self.openapi_client: Optional[CozeOpenAPIClient] = None
        self.websocket_client: Optional[CozeWebSocketClient] = None
        self.is_initialized = False
        
        # 配置对象
        self.credentials: Optional[CozeAPICredentials] = None
        self.agent_config: Optional[CozeAgentConfig] = None
        self.connection_config: Optional[CozeConnectionConfig] = None
        self.flow_config: Optional[CozeFlowConfig] = None
        
        # 回调函数
        self.response_callback: Optional[Callable[[str], None]] = None
        
    async def initialize(self) -> bool:
        """
        初始化Coze集成
        
        Returns:
            初始化是否成功
        """
        try:
            # 加载配置
            await self._load_configurations()
            
            # 验证配置
            if not self._validate_configurations():
                return False
                
            # 初始化客户端
            await self._initialize_clients()
            
            self.is_initialized = True
            logger.info("Coze集成初始化成功")
            return True
            
        except Exception as e:
            logger.error(f"Coze集成初始化失败: {e}")
            return False
            
    async def _load_configurations(self):
        """加载所有配置"""
        # 从环境变量和配置文件合并配置
        self.credentials, self.agent_config = self.config_manager.merge_with_environment()
        
        # 加载连接和流配置
        self.connection_config = self.config_manager.get_connection_config()
        self.flow_config = self.config_manager.get_flow_config()
        
        logger.info("配置加载完成")
        
    def _validate_configurations(self) -> bool:
        """验证配置有效性"""
        if not self.credentials or not self.credentials.api_key:
            logger.error("缺少API凭证")
            return False
            
        if not self.agent_config or not self.agent_config.bot_id:
            logger.error("缺少智能体配置")
            return False
            
        if not self.config_manager.validate_api_key(self.credentials.api_key):
            logger.error("API密钥格式无效")
            return False
            
        if not self.config_manager.validate_bot_id(self.agent_config.bot_id):
            logger.error("智能体ID格式无效")
            return False
            
        return True
        
    async def _initialize_clients(self):
        """初始化客户端"""
        # 初始化OpenAPI客户端
        if self.connection_config.openapi_enabled:
            api_config = CozeAPIConfig(
                api_key=self.credentials.api_key,
                bot_id=self.agent_config.bot_id,
                base_url=self.credentials.base_url,
                timeout=self.connection_config.request_timeout,
                max_retries=self.connection_config.max_retries,
                retry_delay=self.connection_config.retry_delay
            )
            
            self.openapi_client = CozeOpenAPIClient(api_config)
            await self.openapi_client.connect()
            logger.info("OpenAPI客户端初始化完成")
            
        # 初始化WebSocket客户端
        if self.connection_config.websocket_enabled:
            ws_config = WebSocketConfig(
                bot_id=self.agent_config.bot_id,
                base_url=self.connection_config.websocket_url,
                timeout=self.connection_config.connection_timeout,
                max_reconnect_attempts=self.connection_config.max_retries,
                reconnect_delay=self.connection_config.retry_delay
            )
            
            event_handler = VisionWebSocketEventHandler(self.response_callback)
            self.websocket_client = CozeWebSocketClient(ws_config, event_handler)
            
            # 连接WebSocket
            if await self.websocket_client.connect():
                logger.info("WebSocket客户端初始化完成")
            else:
                logger.warning("WebSocket客户端连接失败")
                
    async def shutdown(self):
        """关闭所有连接"""
        try:
            if self.openapi_client:
                await self.openapi_client.disconnect()
                
            if self.websocket_client:
                await self.websocket_client.disconnect()
                
            self.is_initialized = False
            logger.info("Coze集成已关闭")
            
        except Exception as e:
            logger.error(f"关闭Coze集成时出错: {e}")
            
    def set_response_callback(self, callback: Callable[[str], None]):
        """
        设置响应回调函数
        
        Args:
            callback: 回调函数
        """
        self.response_callback = callback
        
        # 更新WebSocket事件处理器
        if self.websocket_client and self.websocket_client.event_handler:
            self.websocket_client.event_handler.response_callback = callback
            
    async def process_vision_query(self, request: VisionQueryRequest) -> VisionQueryResponse:
        """
        处理视觉查询请求
        
        Args:
            request: 视觉查询请求
            
        Returns:
            视觉查询响应
        """
        if not self.is_initialized:
            raise CozeIntegrationError("Coze集成未初始化")
            
        start_time = datetime.now()
        
        try:
            if request.use_websocket and self.websocket_client and self.websocket_client.is_connected:
                # 使用WebSocket方式
                response = await self._process_websocket_query(request)
                response.method_used = "websocket"
            elif self.openapi_client:
                # 使用OpenAPI方式
                response = await self._process_openapi_query(request)
                response.method_used = "openapi"
            else:
                raise CozeIntegrationError("没有可用的客户端")
                
            # 计算处理时间
            processing_time = int((datetime.now() - start_time).total_seconds() * 1000)
            response.processing_time_ms = processing_time
            response.timestamp = datetime.now()
            
            return response
            
        except Exception as e:
            logger.error(f"处理视觉查询失败: {e}")
            
            return VisionQueryResponse(
                query_id=request.query_id,
                response_text="抱歉，处理您的请求时出现了错误。",
                confidence_score=0.0,
                processing_time_ms=int((datetime.now() - start_time).total_seconds() * 1000),
                model_used="error",
                method_used="error",
                timestamp=datetime.now(),
                error=str(e)
            )
            
    async def _process_openapi_query(self, request: VisionQueryRequest) -> VisionQueryResponse:
        """使用OpenAPI处理查询"""
        try:
            chat_response = await self.openapi_client.process_vision_query(
                image_data=request.image_data,
                audio_data=request.audio_data,
                query=request.user_query
            )
            
            return VisionQueryResponse(
                query_id=request.query_id,
                response_text=chat_response.response_text,
                confidence_score=chat_response.confidence_score,
                processing_time_ms=chat_response.processing_time_ms,
                model_used=chat_response.model_used,
                method_used="openapi",
                timestamp=chat_response.timestamp
            )
            
        except CozeAPIError as e:
            raise CozeIntegrationError(f"OpenAPI查询失败: {e}")
            
    async def _process_websocket_query(self, request: VisionQueryRequest) -> VisionQueryResponse:
        """使用WebSocket处理查询"""
        try:
            # 首先上传图像文件（通过OpenAPI）
            if not self.openapi_client:
                raise CozeIntegrationError("需要OpenAPI客户端来上传文件")
                
            image_temp_path = self.openapi_client.save_image_to_temp(request.image_data)
            
            try:
                image_result = await self.openapi_client.upload_file(image_temp_path, "image")
                image_file_id = image_result.file_id
                
                # 通过WebSocket发送多模态查询
                await self.websocket_client.send_multimodal_query(
                    image_file_id=image_file_id,
                    audio_data=request.audio_data
                )
                
                # 等待响应（这里需要实现响应等待机制）
                response_text = await self._wait_for_websocket_response(request.timeout)
                
                return VisionQueryResponse(
                    query_id=request.query_id,
                    response_text=response_text,
                    confidence_score=0.8,  # 默认置信度
                    processing_time_ms=0,  # 将在上层计算
                    model_used="coze",
                    method_used="websocket",
                    timestamp=datetime.now()
                )
                
            finally:
                # 清理临时文件
                self.openapi_client.cleanup_temp_file(image_temp_path)
                
        except Exception as e:
            raise CozeIntegrationError(f"WebSocket查询失败: {e}")
            
    async def _wait_for_websocket_response(self, timeout: int) -> str:
        """
        等待WebSocket响应
        
        Args:
            timeout: 超时时间（秒）
            
        Returns:
            响应文本
        """
        # 这里需要实现一个响应等待机制
        # 可以使用asyncio.Event或者队列来实现
        response_event = asyncio.Event()
        response_text = ""
        
        def response_handler(text: str):
            nonlocal response_text
            response_text = text
            response_event.set()
            
        # 临时设置回调
        original_callback = self.response_callback
        self.set_response_callback(response_handler)
        
        try:
            # 等待响应
            await asyncio.wait_for(response_event.wait(), timeout=timeout)
            return response_text
        except asyncio.TimeoutError:
            return "响应超时，请稍后重试。"
        finally:
            # 恢复原始回调
            self.set_response_callback(original_callback)
            
    async def test_connection(self) -> Dict[str, bool]:
        """
        测试连接状态
        
        Returns:
            连接状态字典
        """
        results = {
            "openapi": False,
            "websocket": False,
            "overall": False
        }
        
        try:
            # 测试OpenAPI连接
            if self.openapi_client:
                # 可以发送一个简单的测试请求
                results["openapi"] = True
                
            # 测试WebSocket连接
            if self.websocket_client and self.websocket_client.is_connected:
                results["websocket"] = True
                
            results["overall"] = results["openapi"] or results["websocket"]
            
        except Exception as e:
            logger.error(f"连接测试失败: {e}")
            
        return results
        
    def get_status(self) -> Dict[str, Any]:
        """
        获取集成状态
        
        Returns:
            状态信息字典
        """
        return {
            "initialized": self.is_initialized,
            "openapi_available": self.openapi_client is not None,
            "websocket_available": self.websocket_client is not None,
            "websocket_connected": (
                self.websocket_client.is_connected 
                if self.websocket_client else False
            ),
            "bot_id": self.agent_config.bot_id if self.agent_config else None,
            "model_name": (
                self.agent_config.model_config.model_name 
                if self.agent_config and self.agent_config.model_config else None
            )
        }
        
    async def setup_from_environment(self) -> bool:
        """
        从环境变量快速设置
        
        Returns:
            设置是否成功
        """
        try:
            env_config = self.config_manager.get_environment_config()
            
            if not env_config['COZE_API_KEY'] or not env_config['COZE_BOT_ID']:
                logger.error("环境变量中缺少必要的配置")
                return False
                
            # 设置默认智能体
            success = self.config_manager.setup_default_agent(
                bot_id=env_config['COZE_BOT_ID'],
                api_key=env_config['COZE_API_KEY']
            )
            
            if success:
                # 重新初始化
                return await self.initialize()
            else:
                return False
                
        except Exception as e:
            logger.error(f"从环境变量设置失败: {e}")
            return False


# 便利函数
async def create_coze_integration(config_dir: Optional[str] = None) -> CozeIntegration:
    """
    创建并初始化Coze集成实例
    
    Args:
        config_dir: 配置文件目录
        
    Returns:
        初始化后的Coze集成实例
    """
    integration = CozeIntegration(config_dir)
    
    if await integration.initialize():
        return integration
    else:
        # 尝试从环境变量设置
        if await integration.setup_from_environment():
            return integration
        else:
            raise CozeIntegrationError("无法初始化Coze集成")


def convert_image_array_to_bytes(image_array: np.ndarray, format: str = "JPEG") -> bytes:
    """
    将numpy图像数组转换为字节数据
    
    Args:
        image_array: 图像数组
        format: 图像格式
        
    Returns:
        图像字节数据
    """
    try:
        from PIL import Image
        import io
        
        # 确保数据类型正确
        if image_array.dtype != np.uint8:
            image_array = (image_array * 255).astype(np.uint8)
            
        # 转换为PIL图像
        if len(image_array.shape) == 3:
            if image_array.shape[2] == 3:
                pil_image = Image.fromarray(image_array, 'RGB')
            elif image_array.shape[2] == 4:
                pil_image = Image.fromarray(image_array, 'RGBA')
            else:
                raise ValueError(f"不支持的图像通道数: {image_array.shape[2]}")
        else:
            pil_image = Image.fromarray(image_array, 'L')
            
        # 转换为字节数据
        buffer = io.BytesIO()
        pil_image.save(buffer, format=format)
        return buffer.getvalue()
        
    except Exception as e:
        logger.error(f"图像转换失败: {e}")
        raise